import pandas as pd
import tushare as ts
import config
import datetime


def getdata(trade_date, det):
    pro = ts.pro_api()

    # (close-open)/open<det，导出文档名
    pfile = trade_date + 'positive.xlsx'
    # (close-open)/open>-det，导出文档名
    nfile = trade_date + 'negative.xlsx'

    positive = pd.DataFrame(columns=['symbol', 'name', 'open', 'close', 'trade_date'])
    negative = pd.DataFrame(columns=['symbol', 'name', 'open', 'close', 'trade_date'])
    positive['symbol'].astype(str)
    negative['symbol'].astype(str)

    # 选取主板股票列表，排除创业板，科创板
    stocklist = pro.stock_basic(exchange='', list_status='L', market='主板', fields='ts_code,symbol,name')

    dtend = datetime.datetime.strptime(trade_date, '%Y%m%d').date()
    dtbegin = dtend - datetime.timedelta(days=30)
    dtbeginstr = dtbegin.strftime('%Y%m%d')

    print(stocklist)
    stocks = stocklist
    codes = ''
    num = 0

    for stock in stocklist.itertuples():
        ts_code = getattr(stock, 'ts_code')
        codes += ts_code
        num = num + 1
        if num == 50:
            # 批量获取行情数据
            datas = pro.daily(ts_code=codes, trade_date=trade_date)

            codes = ''
            num = 0
            # 行情数据过滤
            for df in datas.itertuples():
                ts_code = getattr(df, 'ts_code')
                open = getattr(df, 'open')
                close = getattr(df, 'close')
                trade_date = getattr(df, 'trade_date')

                # 以下3行代码，用于过滤今日收盘价大于昨天收盘价，如果不需要可以注释掉或者删掉
                pre_close = getattr(df, 'pre_close')
                if close <= pre_close:
                    continue
                ###################################################

                #根据收盘价来折算筛选比例，收盘价越小比例越小
                k = 1
                if close < 10:
                    k = 0.2
                elif close < 20:
                    k = 0.4
                elif close < 30:
                    k = 0.5

                # 选取收盘价高于开盘价，并且满足det的数据
                if (close - open) / open < det*k and close - open > 0:
                    #计算10日均线，并过滤收盘价低于10日均线的数据
                    temp = pro.daily(ts_code=ts_code, start_date=dtbeginstr, end_date=trade_date)
                    mean10 = temp.close.rolling(10).mean()
                    mean10 = mean10.dropna()
                    if len(mean10.values) > 0:
                        fmean10 = mean10.values[0]
                        if close < fmean10:
                            continue

                    name = stocklist[stocklist.ts_code == ts_code]['name'].values[0]
                    symbol = stocklist[stocklist.ts_code == ts_code]['symbol'].values[0]
                    serials = pd.Series(
                        {'symbol': symbol, 'name': name, 'open': open, 'close': close, 'trade_date': trade_date},
                        name='')
                    positive = positive.append(serials)

                # 选取收盘价低于开盘价，并且满足det的数据
                if (close - open) / open > -det*k and close - open < 0:
                    #计算10日均线，并过滤收盘价低于10日均线的数据
                    temp = pro.daily(ts_code=ts_code, start_date=dtbeginstr, end_date=trade_date)
                    mean10 = temp.close.rolling(10).mean()
                    mean10 = mean10.dropna()
                    if len(mean10.values) > 0:
                        fmean10 = mean10.values[0]
                        if close < fmean10:
                            continue

                    name = stocklist[stocklist.ts_code == ts_code]['name'].values[0]
                    symbol = stocklist[stocklist.ts_code == ts_code]['symbol'].values[0]
                    serials = pd.Series(
                        {'symbol': symbol, 'name': name, 'open': open, 'close': close, 'trade_date': trade_date},
                        name='')
                    negative = negative.append(serials)
        else:
            codes += ','

    # 将数据保存到excel文件
    positive.columns = ['股票编码', '股票名称', '开盘价', '收盘价', '时间']
    negative.columns = ['股票编码', '股票名称', '开盘价', '收盘价', '时间']
    positive.to_excel(pfile, encoding='gbk')
    negative.to_excel(nfile, encoding='gbk')

if __name__ == '__main__':
    # tushare注册后得到的token
    ts.set_token(config.token)
    # 交易日期
    trade_date = config.trade_date
    # 十字星范围，调整这个参数即可
    det = config.det
    # 获取数据
    getdata(trade_date, det)